import websocket
import json
import base64
from PIL import Image
import io

DETECTION_URL = "ws://localhost:5000"
IMAGE_PATH = "D:\\yolo\\yolov5-master\\yolov5-master\\data\\images\\22.jpg"

def on_message(ws, message):
    """Handle the message received from WebSocket server."""
    print("Received response:")
    print(message)

def on_error(ws, error):
    """Handle errors during WebSocket connection."""
    print(f"Error: {error}")

def on_close(ws, close_status_code, close_msg):
    """Handle WebSocket connection closure."""
    print("Closed connection")

def on_open(ws):
    """Handle WebSocket connection open."""
    # Read image and convert to base64
    with open(IMAGE_PATH, "rb") as image_file:
        image_data = image_file.read()
        base64_image = base64.b64encode(image_data).decode("utf-8")

    # Send image data to the WebSocket server
    data = {
        "image": base64_image,
        "model": "yolov5s"
    }

    # Send data over WebSocket
    ws.send(json.dumps(data))
    print("Image sent to the server.")

# WebSocket client connection setup
ws = websocket.WebSocketApp(DETECTION_URL,
                            on_message=on_message,
                            on_error=on_error,
                            on_close=on_close)
ws.on_open = on_open

# Run the WebSocket client
ws.run_forever()
